package util

import (
	"log/slog"
	"sync"
)

// Task 定义工作池任务接口
type Task interface {
	Execute()
}

// WorkerPool 工作池实现，用于控制并发
type WorkerPool struct {
	workerCount int
	taskQueue   chan Task
	wg          sync.WaitGroup
	logger      *slog.Logger
	shutdown    chan struct{}
}

// NewWorkerPool 创建新的工作池
func NewWorkerPool(workerCount, queueSize int, logger *slog.Logger) *WorkerPool {
	pool := &WorkerPool{
		workerCount: workerCount,
		taskQueue:   make(chan Task, queueSize),
		logger:      logger,
		shutdown:    make(chan struct{}),
	}

	// 启动工作者
	pool.startWorkers()

	return pool
}

// 启动工作者协程
func (p *WorkerPool) startWorkers() {
	for i := 0; i < p.workerCount; i++ {
		p.wg.Add(1)
		go func(workerID int) {
			defer p.wg.Done()
			p.logger.Debug("Worker started", "worker_id", workerID)

			for {
				select {
				case task, ok := <-p.taskQueue:
					if !ok {
						p.logger.Debug("Worker exiting", "worker_id", workerID)
						return
					}
					task.Execute()
				case <-p.shutdown:
					p.logger.Debug("Worker exiting due to shutdown", "worker_id", workerID)
					return
				}
			}
		}(i)
	}
}

// Submit 提交任务到工作池
func (p *WorkerPool) Submit(task Task) bool {
	select {
	case p.taskQueue <- task:
		return true
	default:
		// 队列满，返回失败
		return false
	}
}

// Shutdown 关闭工作池
func (p *WorkerPool) Shutdown() {
	close(p.shutdown)
	close(p.taskQueue)
	p.wg.Wait()
	p.logger.Debug("Worker pool shutdown complete")
}

// PendingTasks 返回等待中的任务数
func (p *WorkerPool) PendingTasks() int {
	return len(p.taskQueue)
}
